/*-
 * Copyright (C) 2018-2020, Shanghai Trusted Industry Control Platform Co., Ltd.
 * Created by 周正辉 (Zhenghui Zhou) <zhouzhenghui@ticpsh.com>, 赵楷文 (Kaiwen
 * Zhao) <zhaokaiwen@ticpsh.com>
 */

#ifndef TESTTARGETVISITOR_H
#define TESTTARGETVISITOR_H

#include <clang/AST/ASTConsumer.h>
#include <clang/AST/Mangle.h>
#include <clang/AST/RecursiveASTVisitor.h>
#include <clang/Frontend/CompilerInstance.h>
#include <clang/Lex/Preprocessor.h>
#include <clang/Lex/Token.h>
#include <llvm/Support/raw_os_ostream.h>

#include <boost/filesystem.hpp>
#include <map>
#include <set>
#include <string>
#include <vector>

#include "share/analyze_setting.h"
#include "share/fixed_address.h"
#include "testhelper.h"

namespace rocket {
namespace testif {

enum SCOPEKIND {
  SCOPE_RECORD,     // class, union, struct
  SCOPE_NAMESPACE,  // namepsace
  SCOPE_FUNCTION,   // function
  SCOPE_LINKAGE,

  SCOPE_TEMPLATED,

  SCOPEKUNKNOW,
};
const std::string SCOPEKIND_STR[] = {"RECORD", "NAMESPACE", "FUNCTION",
                                     "LINKAGE", "TEMPLATED"};
std::string getScopeTypeAsString(const SCOPEKIND &st);

struct TargetDecl {
  TargetDecl() {}
};

struct TargetFunction : public TargetDecl {
  clang::FunctionDecl *declptr = nullptr;
  std::set<clang::CallExpr *> stubs;
  std::set<clang::VarDecl *> globals;
  std::set<clang::VarDecl *> members;
  std::set<clang::FunctionDecl *> funcs;
  std::vector<clang::QualType> types;

  // 局部静态变量
  std::vector<clang::VarDecl *> statics;

  // 异常
  std::vector<clang::CXXThrowExpr *> throwexprs;

  // 成员语句
  std::vector<clang::MemberExpr *> memberexprs;

  // 绝对地址
  std::map<std::string, std::vector<share::FixedAddress>> fixed_addresses;

  // 构造函数作为可选桩函数
  std::set<clang::CXXConstructExpr *> ctors;

  // 模板函数的特化
  std::map<clang::FunctionDecl *, TargetFunction> specfs;

  std::map<clang::VarDecl *, clang::DeclStmt *> locals_decls;

  TargetFunction() {}
  TargetFunction(clang::FunctionDecl *decl_) : declptr(decl_) {}
  TargetFunction(clang::FunctionTemplateDecl *decl_);
};

struct TargetGlobalVar : public TargetDecl {
  clang::VarDecl *declptr = nullptr;

  TargetGlobalVar() {}
  TargetGlobalVar(clang::VarDecl *decl_) : declptr(decl_) {}
};

/**
 * @brief 用于提取待测函数和全局变量
 */
class TestTopLevelVisitor
    : public share::FixedAddressVisitor<TestTopLevelVisitor> {
 public:
  std::map<clang::FunctionDecl *, TargetFunction> &functions_;  // 待测函数
  std::map<const clang::FunctionDecl *, FuncDeclInfo>
      &functiondecls_;  // 所有函数声明信息
  std::map<clang::VarDecl *, TargetGlobalVar> &globals_;  // 全局变量
  const share::RocketAnalyzeSetting &setting;
  std::set<const clang::Type *> types;  // 额外要求解析的类型

  /**
   * @brief 是否进行完全解析
   */
  enum class AnalyzeLevel : unsigned {
    AL_Local = 0,  // 仅本地文件解析
    AL_Used = 1,   // 所有被本地文件中函数直接/间接使用
    AL_Full = 2,   // 该编译单元内所有的函数
  };

  AnalyzeLevel analyze_level = AnalyzeLevel::AL_Local;

 public:
  TestTopLevelVisitor(
      std::map<clang::FunctionDecl *, TargetFunction> &functions,
      std::map<const clang::FunctionDecl *, FuncDeclInfo> &functiondecls,
      std::map<clang::VarDecl *, TargetGlobalVar> &globals,
      clang::ASTContext &ast, const share::RocketAnalyzeSetting &setting,
      unsigned analyze_level)
      : FixedAddressVisitor(ast),
        functions_(functions),
        functiondecls_(functiondecls),
        globals_(globals),
        setting(setting) {
    switch (analyze_level) {
      default:
      case 0: {
        this->analyze_level = AnalyzeLevel::AL_Local;
        break;
      }
      case 1: {
        this->analyze_level = AnalyzeLevel::AL_Used;
        break;
      }
      case 2: {
        this->analyze_level = AnalyzeLevel::AL_Full;
        break;
      }
    }
  }

  void addOverriding(const clang::CXXMethodDecl *ocerridden,
                     const clang::CXXMethodDecl *overridding);
  std::map<const clang::FunctionDecl *, FuncDeclInfo>::iterator FuncDeclAction(
      const clang::FunctionDecl *fd);

  // get target functions
  virtual bool VisitFunctionDecl(clang::FunctionDecl *fd);
  virtual bool VisitFunctionTemplateDecl(clang::FunctionTemplateDecl *ftd);

  // get toplevel globals
  bool VisitVarDecl(clang::VarDecl *nd);

  // FIXED: 复杂项目中直接访问类特化会导致因数据量过多从而段错误
  // get spec record method
  bool VisitClassTemplateDecl(clang::ClassTemplateDecl *ctsd);

  bool VisitEnumDecl(clang::EnumDecl *ed);

  bool TraverseDecl(clang::Decl *dl);

  bool addTargetFunction(clang::FunctionDecl *fd);
  bool addTargetFunction(clang::FunctionTemplateDecl *ftd);
  bool analyzeTargetFunction(TargetFunction &tf);
  bool analyzeStubFunction(TargetFunction &tf, clang::FunctionDecl *fd);
};

/**
 * @brief 用于提取待某一结构中使用的全局变量
 */
class InnerGlobalVisitor
    : public share::FixedAddressVisitor<InnerGlobalVisitor> {
 public:
  std::set<clang::FunctionDecl *> functions_;
  std::set<clang::VarDecl *> globals_;
  std::set<clang::UserDefinedLiteral *> udliteral_;

  std::vector<clang::MemberExpr *> member_exprs;

 public:
  InnerGlobalVisitor(clang::ASTContext &ast) : FixedAddressVisitor(ast) {}
  bool VisitDeclRefExpr(clang::DeclRefExpr *dre);
  bool VisitCXXConstructExpr(clang::CXXConstructExpr *ctor);
  bool VisitUserDefinedLiteral(clang::UserDefinedLiteral *udl);

  bool VisitMemberExpr(clang::MemberExpr *me);
};

/**
 * @brief 用于解析待测函数
 */
class TestIfVisitor : public share::FixedAddressVisitor<TestIfVisitor> {
 private:
  TargetFunction &func_;
  const share::RocketAnalyzeSetting &setting;

 public:
  TestIfVisitor(clang::ASTContext &ast, TargetFunction &func,
                const share::RocketAnalyzeSetting &setting)
      : FixedAddressVisitor(ast), func_(func), setting(setting) {}

 public:
  bool VisitCallExpr(clang::CallExpr *ce);
  bool VisitDeclRefExpr(clang::DeclRefExpr *e);
  bool VisitCXXConstructExpr(clang::CXXConstructExpr *cxxce);
  bool VisitDecl(clang::Decl *decl);
  bool VisitCastExpr(clang::CastExpr *expr);
  bool VisitCXXThrowExpr(clang::CXXThrowExpr *te);

  bool VisitMemberExpr(clang::MemberExpr *me);

  bool VisitCXXForRangeStmt(clang::CXXForRangeStmt *stmt);

  bool VisitCXXDefaultArgExpr(clang::CXXDefaultArgExpr *expr);

  bool VisitMaterializeTemporaryExpr(clang::MaterializeTemporaryExpr *tExpr);

  bool VisitCXXNewExpr(clang::CXXNewExpr *ne);
  bool VisitCXXDeleteExpr(clang::CXXDeleteExpr *de);

  bool VisitDeclStmt(clang::DeclStmt *ds);

  bool analyzeStubFunction(clang::FunctionDecl *fd);

 private:
  std::set<clang::Expr *> traversedExpr;
  std::set<clang::Stmt *> traversedStmt;
  std::set<clang::Decl *> traversedDecl;
  bool dealUntraversedNode(clang::Expr *expr);
};

}  // namespace testif
}  // namespace rocket

#endif  // TESTTARGETVISITOR_H
